__author__ = 'xinrong'

'''
Created on Oct 6, 2010

@author: Peter
'''
from numpy import *
import matplotlib.pyplot as plt


decisionNode = dict(boxstyle="sawtooth", fc="0.8")
leafNode = dict(boxstyle="round4", fc="0.8")
arrow_args = dict(arrowstyle="<-")


class Tree(object):
    def __init__(self, value=None, left=None, right=None, split=True, sign=-1, parent=None):
        self.value = value
        self.left = left
        self.right = right
        self.split = split
        self.parent = parent
        self.ck = sign

    def set_sign(self, s):
        self.sig = s

    def get_sign(self):
        return self.sig

    def set_value(self, v):
        self.value = v

    def get_value(self):
        return self.value

    def get_left(self):
        return self.left

    def set_left(self, l):
        self.left = l

    def get_right(self):
        return self.right

    def set_right(self, r):
        self.right = r

    def get_parent(self):
        return self.parent

    def print_tree(self, off):
        ss = '    '
        print '%s %s' % (ss * off, self.value)
        if self.left is not None:
            self.left.print_tree(off - 4)
        if self.right is not None:
            self.right.print_tree(off + 4)
            # print("\n")

    def print_shape(self, off):
        ss = ' ' * 8
        sp = '$'

        full_num = 2
        current = 0
        lst = [self, sp]
        print '%s' % ss * (off - int(math.log(full_num))),
        while len(lst) > 0:

            flag = True
            for k in lst:
                if k is not ss and k is not sp:
                    flag = False
                    break

            if flag:
                break

            em = lst.pop(0)
            flag = 'x'
            if em.split is True:
                flag = 'x'
            else:
                flag = 'y'

            if isinstance(em, Tree):
                print '%s%s %s' % (em.value, flag, ss),

                if em.get_left() is None:
                    lst.append(ss)
                else:
                    lst.append(em.get_left())
                if em.get_right() is None:
                    lst.append(ss)
                else:
                    lst.append(em.get_right())
                current += 2
            elif em == sp:
                print ''
                print '%s' % ss * (off - int(math.log(full_num))),
            else:
                print '%s %s' % (ss * off, ss),
                lst.append(ss)
                lst.append(ss)
                current += 2
            if current == full_num:
                full_num *= 2
                current = 0
                lst.append(sp)


def find_mid(arr, pivot):
    if pivot is True:
        temp = []
        for ix in range(len(arr)):
            temp_len = len(temp)
            pos = 0
            while pos < temp_len and arr[ix][0] > temp[pos][0]:
                pos += 1
            temp.insert(pos, arr[ix])
        arr = temp
    else:
        temp = []
        for iy in range(len(arr)):
            temp_len = len(temp)
            pos = 0
            while pos < temp_len and arr[iy][1] > temp[pos][1]:
                pos += 1
            temp.insert(pos, arr[iy])
        arr = temp
    return arr[len(arr) / 2], arr


def make_tree(arr, pivot):
    temp = Tree()
    if len(arr) > 1:
        temp.value, temp_arr = find_mid(arr, pivot)
        if len(temp_arr) / 2 >= 1:
            temp.set_left(make_tree(temp_arr[0: len(temp_arr) / 2], not pivot))
            temp.left.parent = temp
        if len(temp_arr) / 2 + 1 < len(arr):
            temp.set_right(make_tree(temp_arr[len(temp_arr) / 2 + 1:len(temp_arr)], not pivot))
            temp.right.parent = temp
    elif len(arr) == 1:
        temp.set_value(arr[0])
    else:
        return None
    temp.split = pivot
    return temp


def heappush(heap, item, dest):
    """Push item onto heap, maintaining the heap invariant."""
    heap.append(item)
    _siftdown(heap, 0, len(heap) - 1, dest)


def _siftdown(heap, startpos, pos, dest):
    newitem = heap[pos]
    # Follow the path to the root, moving parents down until finding a place
    # newitem fits.
    while pos > startpos:
        parent = heap[pos - 1]
        if distance_cmp_lt(newitem, parent, dest):
            heap[pos] = parent
            pos -= 1
            continue
        break
    heap[pos] = newitem


def distance_cmp_lt(first, second, center=[0, 0]):
    # Use __lt__ if available; otherwise, try __le__.
    # In Py3.x, only __lt__ will be called.

    dx = (first.value[0] - center[0]) ** 2 + (first.value[1] - center[1]) ** 2
    dy = (second.value[0] - center[0]) ** 2 + (second.value[1] - center[1]) ** 2
    # return (dx < dy) if hasattr(dx, '__lt__') else (not dy <= dx)
    return dx < dy  # asc


def deleteNode(t, rn):
    if t is None or t.ck == rn:
        return
    else:
        t.ck = rn

        # p = t.parent
        # if p is None or p.ck == rn:
        # return
        # if p.split is True:
        # if p.value[0] > t.value[0]:
        #         p.left.ck = rn
        #     else:
        #         p.right.ck = rn
        # else:
        #     if p.value[1] > t.value[1]:
        #         p.left.ck = rn
        #     else:
        #         p.right.ck = rn


def otherbranch(t):
    p = t.parent
    if p.split is True:
        if p.value[0] > t.value[0]:  # left of parent
            return p.right
        else:
            return p.left  # right of parent
    else:
        if p.value[1] > t.value[1]:  # down of parent
            return p.right
        else:  # up of parent
            return p.left


def addnode(lst, p, elem, n, rn):
    if p is None or p.ck == rn:
        return lst
    else:
        if lst is None:
            lst = []
        else:
            heappush(lst, p, elem)


def judgeCross(l, p, e, n, rn):
    if p is None or p.ck == rn:
        return False
    elif len(l) < n:
        return True
    else:
        d1 = (l[0].value[0] - e[0]) ** 2 + (l[0].value[1] - e[1]) ** 2
        d2 = (l[len(l) - 1].value[0] - e[0]) ** 2 + (l[len(l) - 1].value[1] - e[1]) ** 2
        d2 = d1 if d1 > d2 else d2
        s = (e[1] - p.parent.value[1]) ** 2 if p.parent.split is True else  (e[0] - p.parent.value[0]) ** 2
        return True if s < d2 else False


def findKNN(t, elem, k, rn, nlist=None):
    if nlist is None:
        nlist = []
    if t is None or t.ck == rn:
        return nlist[0:k]
    elif not isinstance(t, Tree):
        raise "not tree "
    else:
        if len(nlist) >= k:
            nlist = nlist[0:k]
        if t.split is True:
            if elem[0] < t.value[0]:  # find left
                if t.left is not None and t.left.ck != rn:
                    return findKNN(t.left, elem, k, rn, nlist)
                else:
                    if judgeCross(nlist, t.right, elem, k, rn):
                        return findKNN(t.right, elem, k, rn, nlist)
                    else:
                        addnode(nlist, t, elem, k, rn)
                        deleteNode(t, rn)
                        return findKNN(t.parent, elem, k, rn, nlist)
            else:  # find right
                if t.right is not None and t.right.ck != rn:
                    return findKNN(t.right, elem, k, rn, nlist)
                else:
                    if judgeCross(nlist, t.left, elem, k, rn):
                        return findKNN(t.left, elem, k, rn, nlist)
                    else:
                        addnode(nlist, t, elem, k, rn)
                        deleteNode(t, rn)
                        return findKNN(t.parent, elem, k, rn, nlist)
        else:
            if elem[1] < t.value[1]:  # find down
                if t.left is not None and t.left.ck != rn:
                    return findKNN(t.left, elem, k, rn, nlist)
                else:
                    if judgeCross(nlist, t.right, elem, k, rn):
                        return findKNN(t.right, elem, k, rn, nlist)
                    else:
                        addnode(nlist, t, elem, k, rn)
                        deleteNode(t, rn)
                        return findKNN(t.parent, elem, k, rn, nlist)
            else:  # find up
                if t.right is not None and t.right.ck != rn:
                    return findKNN(t.right, elem, k, rn, nlist)
                else:
                    if judgeCross(nlist, t.left, elem, k, rn):
                        return findKNN(t.left, elem, k, rn, nlist)
                    else:
                        addnode(nlist, t, elem, k, rn)
                        deleteNode(t, rn)
                        return findKNN(t.parent, elem, k, rn, nlist)


def get_level(node):
    if node is None:
        return 0
    if isinstance(node, Tree):
        left_l = get_level(node.left)
        right_l = get_level(node.right)
        return left_l + 1 if left_l > right_l else right_l + 1
    else:
        return 0


def get_leafs(node):
    if node is None:
        return 0
    elif isinstance(node, Tree):
        if node.left is None and node.right is None:
            return 1
        else:
            left_l = get_level(node.left)
            right_l = get_level(node.right)
            return left_l + right_l
    else:
        return 0


def plotTree(myTree, parentPt, nodeTxt):  # if the first key tells you what feat was split on
    numLeafs = get_leafs(myTree)  #this determines the x width of this tree
    levels = get_level(myTree)
    # firstStr = myTree.keys()[0]     #the text label for this node should be this
    # if x:
    cntrPt = (plotTree.xOff + (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff)
    # else:
    #     cntrPt = (plotTree.xOff - (1.0 + float(numLeafs))/2.0/plotTree.totalW, plotTree.yOff)

    # if x:
    #     cntrPt = (parentPt[0] - 0.13 * 2 ** levels, plotTree.yOff)
    # else:
    #     cntrPt = (parentPt[0] + 0.13 * 2 ** levels, plotTree.yOff)

    plotMidText(cntrPt, parentPt, nodeTxt)
    plotNode(myTree.value, cntrPt, parentPt, decisionNode)
    plotTree.yOff = plotTree.yOff - 1.0 / plotTree.totalD
    if myTree.left is not None:
        if get_level(myTree.left) == 1:
            plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(myTree.left.value))
        else:
            plotTree(myTree.left, cntrPt, 'x' if myTree.split else 'y')
    if myTree.right is not None:
        if get_level(myTree.right) == 1:
            plotTree.xOff = plotTree.xOff + 1.0/plotTree.totalW
            plotMidText((plotTree.xOff, plotTree.yOff), cntrPt, str(myTree.right.value))
        else:
            plotTree(myTree.right, cntrPt, 'x' if myTree.split else 'y')
    plotTree.yOff = plotTree.yOff + 1.0 / plotTree.totalD


# string , end ,start, type
def plotNode(nodeTxt, centerPt, parentPt, nodeType):
    create_plot.ax1.annotate(nodeTxt, xy=parentPt, xycoords='axes fraction',
                            xytext=centerPt, textcoords='axes fraction',
                            va="center", ha="center", bbox=nodeType, arrowprops=arrow_args)


def plotMidText(cntrPt, parentPt, txtString):
    xMid = (parentPt[0] - cntrPt[0]) / 2.0 + cntrPt[0]
    yMid = (parentPt[1] - cntrPt[1]) / 2.0 + cntrPt[1]
    create_plot.ax1.text(xMid, yMid, txtString, va="center", ha="center", rotation=30)


def create_plot(tree_in):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    create_plot.ax1 = plt.subplot(111, frameon=False, **axprops)  # no ticks

    # plotTree.totalW = float(get_leafs(inTree))
    plotTree.totalW = float(get_leafs(tree_in))
    plotTree.totalD = float(get_level(tree_in))
    plotTree.xOff = - 0.5 / plotTree.totalW;
    plotTree.yOff = 1.0;
    plotTree(tree_in, (0.5, 1.0), '')
    plt.show()


n = 20  # number of points to create
markers = []
colors = []
for i in range(n):
    e1 = random.randint(200)
    e2 = random.randint(200)
    markers.append([e1 - 100, e2 - 100])

print markers

tr = make_tree(markers, True)
tr.print_shape(5)

print '\nstart find'
kid = [3, 1]
print kid
rn = random.randint(1, 44)
print ''
l1 = findKNN(tr, kid, 10, rn)
for i in l1:
    v = i.value
    d = (v[0] - kid[0]) ** 2 + (v[1] - kid[1]) ** 2
    print i.value, d

create_plot(tr)




# rn = random.randint(1, 44)
# print ''
# l2 = findKNN(tr, kid, 2, rn)
# for i in l2:
# v = i.value
# d = (v[0]-kid[0])**2+(v[1]-kid[1])**2
#     print i.value, d
#
# rn = random.randint(1, 44)
# print ''
# l3 = findKNN(tr, kid, 3, rn)
# for i in l3:
#     v = i.value
#     d = (v[0]-kid[0])**2+(v[1]-kid[1])**2
#     print i.value, d
#
# rn = random.randint(1, 44)
# print rn
# l4 = findKNN(tr, kid, 4, rn)
# for i in l4:
#     v = i.value
#     d = (v[0]-kid[0])**2+(v[1]-kid[1])**2
#     print i.value, d
